import numpy as np
import scipy.io.wavfile
import io
import pydub
import soundfile as sf
from scipy.signal import resample
import librosa
import sounddevice as sd
from scipy.spatial.distance import cosine


def segment_to_numpy(sound):
    channel_sounds = sound.split_to_mono()
    samples = [s.get_array_of_samples() for s in channel_sounds]
    fp_arr = np.array(samples).T.astype(np.float32)
    fp_arr /= np.iinfo(samples[0].typecode).max
    return fp_arr


def numpy_to_segment(fp_arr, sample_rate):
    wav_io = io.BytesIO()
    scipy.io.wavfile.write(wav_io, sample_rate, fp_arr)
    wav_io.seek(0)
    sound = pydub.AudioSegment.from_wav(wav_io)
    return sound


def load_audio_as_array(raw_audio_path, max_duration, sample_rate_, compress_rate, raw=None):
    if raw is None:
        raw, raw_sr = sf.read(raw_audio_path)
    else:
        raw_sr = sample_rate_

    raw_length = raw.shape[0]

    seg = pydub.AudioSegment.silent(duration=int(max_duration * 1000), frame_rate=sample_rate_)

    comp_length = int(seg.duration_seconds * seg.frame_rate / compress_rate)

    seg1 = pydub.AudioSegment.from_wav(raw_audio_path)[:int(max_duration * 1000)]
    seg1.set_frame_rate(seg.frame_rate)
    seg1.set_sample_width(seg.sample_width)

    arr = segment_to_numpy(seg)
    arr1 = segment_to_numpy(seg1)
    print(max_duration, sample_rate_, arr.shape[0], arr1.shape[0])
    if arr.shape[0] > arr1.shape[0]:
        arr1 = arr1[:arr.shape[0], ...]
        seg1 = numpy_to_segment(arr1, sample_rate_)
    if seg.duration_seconds > seg1.duration_seconds:
        seg = seg.overlay(seg1)
    else:
        seg = seg1.overlay(seg)

    seg = seg.overlay(seg1)
    seg.set_frame_rate(seg1.frame_rate)
    seg.set_sample_width(seg1.sample_width)

    arr = segment_to_numpy(seg)

    dat_0 = resample(arr[:, 0], comp_length)
    if len(arr.shape) == 2 and arr.shape[1] == 2:
        dat_1 = resample(arr[:, 1], comp_length)
    else:
        dat_1 = dat_0

    return np.stack([dat_0, dat_1], axis=-1), raw_length, raw_sr


def audio_similarity(audio1, audio2, hop_length=128):
    _sum = 0
    _range = np.arange(0, audio1.shape[0], hop_length)
    for i in _range:
        if i+100 <= audio1.shape[0]:
            _sum += cosine(audio1[i:i+hop_length, 0], audio2[i:i+hop_length, 0])
            _sum += cosine(audio1[i:i+hop_length, 1], audio2[i:i+hop_length, 1])
    _sum /= _range.shape[0] * 2
    return 1 - _sum


if __name__ == '__main__':
    audio_a, _, _ = load_audio_as_array(r"C:\dev_spa\DMuse\202202c1\clap\Clap 1.wav", 1, 44100, 1)
    audio_b, _, _ = load_audio_as_array(r"C:\dev_spa\DMuse\202202c1\clap\Clap 2.wav", 1, 44100, 1)
    audio_c, _, _ = load_audio_as_array(r"C:\dev_spa\DMuse\Dajun 20201213_01.wav", 1, 44100, 1)

    print(audio_similarity(audio_a, audio_b))
    print(audio_similarity(audio_b, audio_c))
